import os
import uuid
import requests
import json as _json
import base64
import io as _io
import tempfile
from datetime import datetime
import platform
import pathlib
import pandas as pd
import csv
import re
from urllib.parse import urlparse
import threading
import random
from hashlib import md5

import socket

import gradio as gr
from PIL import Image

# 获取环境变量中的路径和 API Key
bozo = os.environ.get("BOZO")
if bozo:
    bozo = bozo + "/Downloads/"
    os.chdir(bozo)

# 获取当天日期，格式化为 "年-月-日" 或 "年月日"（避免特殊字符）
today = datetime.today()
# 可选格式："%Y-%m-%d" 生成 "2024-09-04"；"%Y%m%d" 生成 "20240904"
date_folder = today.strftime("%Y%m%d")  # 示例：20240904

# 构造完整路径（在downloads目录下创建日期文件夹）
target_path = os.path.join("downloads", date_folder)

# 检查路径是否存在，不存在则创建（包括多级目录）
if not os.path.exists(target_path):
    os.makedirs(target_path)  # 自动创建父目录（如downloads不存在也会创建）
    # print(f"已创建文件夹：{target_path}")
else:
    print(f"文件夹已存在：{target_path}")

# # 创建本地缓存目录
# if not os.path.exists("downloads"):
#     os.makedirs("downloads")

# 全局变量用于记录提交次数和key索引
submission_count = 0
current_key_index = 0
key_lock = threading.Lock()  # 线程锁确保key索引安全
gradio_url = "http://0.0.0.0:8003"  # 默认Gradio URL

# 读取API Keys文件
def load_api_keys():
    """从BOZO_key.env文件加载API keys"""
    keys = []
    key_file_path = os.path.join(os.path.dirname(__file__), 'BOZO_key.env')
    
    if not os.path.exists(key_file_path):
        print(f"警告: API key文件 {key_file_path} 不存在")
        return keys
    
    try:
        with open(key_file_path, 'r', encoding='utf-8') as file:
            for line in file:
                line = line.strip()
                # 跳过空行和注释行
                if line and not line.startswith('#'):
                    keys.append(line)
        print(f"成功加载 {len(keys)} 个API keys")
        return keys
    except Exception as e:
        print(f"读取API key文件错误: {e}")
        return keys

# 获取下一个API Key
def get_next_api_key(custom_key=None):
    """获取下一个可用的API key，优先级：自定义输入（自动保存） > 环境变量ORAPI（非空时） > 文件keys"""
    global current_key_index, submission_count
    
    with key_lock:
        submission_count += 1
        print(f"第 {submission_count} 次提交")
        
        # 第一优先级：自定义输入的key（同时保存到文件）
        if custom_key and custom_key.strip():
            custom_key = custom_key.strip()
            print("使用自定义输入的API key")
            # 保存自定义key到文件
            save_key_to_file(custom_key)
            return custom_key
        
        # 第二优先级：环境变量ORAPI（只有非空时才使用）
        env_key = os.environ.get("ORAPI")
        if env_key and env_key.strip():
            print("使用环境变量ORAPI中的API key")
            return env_key.strip()
        
        # 第三优先级：文件中的keys（轮换使用）
        api_keys = load_api_keys()
        if not api_keys:
            raise ValueError("没有可用的API key：请输入自定义key、设置ORAPI环境变量或在BOZO_key.env文件中添加keys")
        
        # 按顺序轮换获取key
        selected_key = api_keys[current_key_index % len(api_keys)]
        print(f"使用文件中第 {current_key_index + 1} 个API key（共 {len(api_keys)} 个）")
        
        # 更新索引（下次提交时使用下一个key）
        current_key_index = (current_key_index + 1) % len(api_keys)
        
        return selected_key

# 百度翻译API配置
BAIDU_APPID = '20230502001663178'
BAIDU_APPKEY = 'Dz2_wutD286PSys6hHLa'

# 百度翻译功能
def make_md5(s, encoding='utf-8'):
    """生成MD5哈希值"""
    return md5(s.encode(encoding)).hexdigest()

def baidu_translate(query, from_lang='zh', to_lang='en'):
    """调用百度翻译API"""
    # 如果输入的是英文，直接返回原文
    if from_lang == 'en' or (from_lang == 'auto' and query.isascii()):
        return query
    
    endpoint = 'http://api.fanyi.baidu.com'
    path = '/api/trans/vip/translate'
    url = endpoint + path
    
    # 生成salt和sign
    salt = random.randint(32768, 65536)
    sign = make_md5(BAIDU_APPID + query + str(salt) + BAIDU_APPKEY)
    
    # 构建请求
    headers = {'Content-Type': 'application/x-www-form-urlencoded'}
    payload = {
        'appid': BAIDU_APPID,
        'q': query,
        'from': from_lang,
        'to': to_lang,
        'salt': salt,
        'sign': sign
    }
    
    try:
        # 发送请求
        r = requests.post(url, params=payload, headers=headers, timeout=100)
        result = r.json()
        
        # 检查是否有翻译结果
        if 'trans_result' in result and result['trans_result']:
            # 合并所有翻译结果
            translated_texts = []
            for item in result['trans_result']:
                translated_texts.append(item['dst'])
            return ' '.join(translated_texts)
        else:
            print(f"百度翻译API响应异常: {result}")
            return query  # 翻译失败时返回原文
            
    except Exception as e:
        print(f"百度翻译失败: {e}")
        return query  # 翻译失败时返回原文

# 保存API Key到文件（带去重功能）
def save_key_to_file(new_key):
    """将新的API key保存到BOZO_key.env文件中，去除重复"""
    key_file_path = os.path.join(os.path.dirname(__file__), 'BOZO_key.env')
    
    # 读取现有的keys
    existing_keys = set()
    comments = []
    
    if os.path.exists(key_file_path):
        try:
            with open(key_file_path, 'r', encoding='utf-8') as file:
                for line in file:
                    line = line.strip()
                    if line.startswith('#') or not line:
                        comments.append(line)
                    else:
                        existing_keys.add(line)
        except Exception as e:
            print(f"读取现有key文件错误: {e}")
    
    # 检查是否重复
    if new_key in existing_keys:
        print(f"API key已存在，跳过添加")
        return
    
    # 添加新key
    existing_keys.add(new_key)
    
    try:
        # 重写文件
        with open(key_file_path, 'w', encoding='utf-8') as file:
            # 先写入注释
            for comment in comments:
                file.write(comment + '\n')
            
            # 再写入所有keys（去重后）
            for key in sorted(existing_keys):  # 排序保持一致性
                file.write(key + '\n')
        
        print(f"成功将新API key添加到文件中")
    except Exception as e:
        print(f"保存API key到文件错误: {e}")

# 更新Gradio URL
def update_gradio_url(local_url, share_url=None):
    """更新全局Gradio URL"""
    global gradio_url
    # 优先使用公共URL，如果没有则使用本地URL
    gradio_url = share_url if share_url else local_url
    print(f"更新Gradio URL为: {gradio_url}")

# 检测可用端口并设置Gradio URL
def setup_gradio_url(port=8003, host="0.0.0.0"):
    """设置并更新Gradio URL"""
    global gradio_url
    
    # 检测端口是否可用
    def is_port_available(port):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            try:
                s.bind((host, port))
                return True
            except:
                return False
    
    # 如果端口不可用，寻找下一个可用端口
    while not is_port_available(port):
        port += 1
        if port > 8004:  # 防止无限循环
            port = 8003
            break
    
    # 设置全局URL
    if host == "0.0.0.0":
        gradio_url = f"http://localhost:{port}"
    else:
        gradio_url = f"http://{host}:{port}"
    
    print(f"设置Gradio URL为: {gradio_url}")
    return port

# # 获取系统默认下载文件夹路径
# def get_default_download_path():
#     if platform.system() == 'Windows':
#         return os.path.join(os.environ['USERPROFILE'], 'Downloads')
#     elif platform.system() == 'Darwin':  # macOS
#         return os.path.join(os.path.expanduser('~'), 'Downloads')
#     else:  # Linux 等其他系统
#         return os.path.join(os.path.expanduser('~'), 'Downloads')


def get_default_download_path():
    # 1. 获取系统默认的Downloads文件夹路径
    default_downloads = os.path.join(os.path.dirname(__file__), 'Downloads')
    
    # 2. 生成当天日期文件夹名称（格式：20240904）
    today = datetime.today()
    date_folder = today.strftime("%Y%m%d")  # 可修改格式，如"%Y-%m-%d"
    
    # 3. 拼接完整路径：默认Downloads文件夹 + 日期子文件夹
    target_path = os.path.join(default_downloads, date_folder)
    
    # 4. 确保文件夹存在（不存在则创建）
    if not os.path.exists(target_path):
        os.makedirs(target_path)
        print(f"已创建日期文件夹：{target_path}")
    else:
        print(f"日期文件夹已存在：{target_path}")
    
    # 5. 返回最终路径（带当天日期的文件夹）
    return target_path


# 读取styles.csv文件
def load_styles():
    """加载样式文件"""
    styles = []
    csv_path = os.path.join(os.path.dirname(__file__), 'styles.csv')
    
    if not os.path.exists(csv_path):
        print(f"警告: 样式文件 {csv_path} 不存在")
        return [], {}
    
    try:
        with open(csv_path, 'r', encoding='utf-8') as file:
            reader = csv.DictReader(file)
            styles_dict = {}
            names = []
            for row in reader:
                name = row.get('name', '').strip()
                prompt = row.get('prompt', '').strip()
                remark = row.get('remark', '').strip()
                
                if name:
                    names.append(name)
                    styles_dict[name] = {
                        'prompt': prompt,
                        'remark': remark
                    }
            
        return names, styles_dict
    except Exception as e:
        print(f"读取样式文件错误: {e}")
        return [], {}

# 从remark中提取图片URL
def extract_image_url(remark):
    """从remark字段中提取图片URL"""
    if not remark:
        return None
    
    # 使用正则表达式匹配HTTP/HTTPS URL
    url_pattern = r'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\(\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+'
    urls = re.findall(url_pattern, remark)
    
    if urls:
        return urls[0]  # 返回第一个找到的URL
    return None

# 从URL下载图片
def download_image_from_url(url):
    """从URL下载图片并返回PIL Image对象"""
    if not url:
        return None
    
    try:
        response = requests.get(url, timeout=100)
        response.raise_for_status()
        
        # 检查是否为图片类型
        content_type = response.headers.get('content-type', '')
        if not content_type.startswith('image/'):
            print(f"警告: URL返回的不是图片类型: {content_type}")
            return None
        
        image = Image.open(_io.BytesIO(response.content))
        return image
    except Exception as e:
        print(f"下载图片失败 {url}: {e}")
        return None

# 效果查看函数
def preview_style(selected_style, styles_dict):
    """预览选中的样式效果"""
    if not selected_style or selected_style not in styles_dict:
        return "", "", None
    
    style_data = styles_dict[selected_style]
    prompt = style_data.get('prompt', '')
    remark = style_data.get('remark', '')
    
    # 提取并下载图片
    image_url = extract_image_url(remark)
    preview_image = None
    if image_url:
        preview_image = download_image_from_url(image_url)
    
    return prompt, remark, preview_image

def query_api_key_usage(custom_key):
    """查询API Key的使用情况"""
    try:
        # 获取要查询的API Key
        query_key = None
        key_source = ""
        
        # 确定使用哪个key进行查询（不更新提交计数）
        if custom_key and custom_key.strip():
            query_key = custom_key.strip()
            key_source = "自定义输入"
            # 保存自定义key到文件（查询时也需要保存）
            save_key_to_file(query_key)
        else:
            env_key = os.environ.get("ORAPI")
            if env_key and env_key.strip():
                query_key = env_key.strip()
                key_source = "环境变量ORAPI"
            else:
                # 使用文件中的第一个key进行查询（不轮换）
                api_keys = load_api_keys()
                if api_keys:
                    query_key = api_keys[0]
                    key_source = f"文件中第1个key（共{len(api_keys)}个）"
        
        if not query_key:
            return "错误: 没有可用的API Key进行查询\n请输入自定义key、设置ORAPI环境变量或在BOZO_key.env文件中添加keys"
        
        # 调用OpenRouter API查询Key使用情况
        response = requests.get(
            url="https://openrouter.ai/api/v1/key",
            headers={
                "Authorization": f"Bearer {query_key}"
            },
            timeout=100
        )
        
        response.raise_for_status()
        result = response.json()
        
        # 格式化显示data部分的响应数据
        if 'data' in result:
            data = result['data']
            
            # 构建显示信息
            info_lines = []
            info_lines.append("=== API Key 使用情况 ===")
            info_lines.append(f"查询的Key来源: {key_source}")
            info_lines.append(f"标签: {data.get('label', 'N/A')}")
            info_lines.append(f"已使用信用额度: {data.get('usage', 0)}")
            
            limit = data.get('limit')
            if limit is None:
                info_lines.append("信用额度限制: 无限制")
            else:
                info_lines.append(f"信用额度限制: {limit}")
                if limit > 0:
                    remaining = limit - data.get('usage', 0)
                    info_lines.append(f"剩余信用额度: {remaining}")
            
            is_free_tier = data.get('is_free_tier', False)
            info_lines.append(f"免费账户: {'是' if is_free_tier else '否'}")
            
            # 添加原始JSON数据
            info_lines.append("\n=== 原始响应数据 ===")
            info_lines.append(_json.dumps(data, indent=2, ensure_ascii=False))
            
            return "\n".join(info_lines)
        else:
            return f"API响应格式异常: {_json.dumps(result, indent=2, ensure_ascii=False)}"
            
    except requests.exceptions.RequestException as e:
        return f"网络请求错误: {str(e)}"
    except Exception as e:
        return f"查询API Key使用情况失败: {str(e)}"

def generate_by_nano_banana(prompt, image1, image2, image3, api_key):
    # 获取API Key（按优先级：自定义 > 环境变量 > 文件）
    openrouter_key = get_next_api_key(api_key)
    if not openrouter_key:
        raise ValueError("没有可用的API Key")

    # 保存原始中文提示词
    original_chinese_prompt = prompt
    
    # 检测是否为中文并进行翻译
    if prompt and not prompt.isascii() and any('\u4e00' <= char <= '\u9fff' for char in prompt):
        # 包含中文字符，进行翻译
        translated_prompt = baidu_translate(prompt, from_lang='zh', to_lang='en')
        print(f"中文提示词: {original_chinese_prompt}")
        print(f"英文提示词: {translated_prompt}")
        final_prompt = translated_prompt
    else:
        # 英文或其他语言，直接使用
        print(f"英文提示词: {prompt}")
        final_prompt = prompt
        original_chinese_prompt = ""  # 没有中文提示词
    
    content_info = [{"type": "text", "text": final_prompt}]
    cache_id = f'{uuid.uuid4()}'
    
    # 存储图片缓存路径用于后续保存
    image_cache_paths = []

    # 处理图片并转换为base64
    for i, image in enumerate([image1, image2, image3]):
        if image:
            # 保存图片到本地缓存
            image_cache_path = f'downloads/{cache_id}_{i}.webp'
            image.save(image_cache_path, format='WebP', quality=95)
            image_cache_paths.append(image_cache_path)  # 记录缓存路径
            
            # 转换为base64格式（API可直接识别）
            with open(image_cache_path, "rb") as img_file:
                base64_str = base64.b64encode(img_file.read()).decode('utf-8')
            
            content_info.append({
                "type": "image_url", 
                "image_url": {"url": f"data:image/webp;base64,{base64_str}"}
            })

    messages = [{"role": "user", "content": content_info}]

    headers = {
        "Authorization": f"Bearer {openrouter_key}",
        "Content-Type": "application/json",
        "HTTP-Referer": gradio_url,  # 使用当前Gradio网页URL
        "X-Title": "BOZO - OpenRouter.ai", 
    }

    data = {
        "model": "google/gemini-2.5-flash-image-preview",
        "messages": messages,
        "temperature": 0.7,
        "max_tokens": 1000
    }

    response = requests.post(
        "https://openrouter.ai/api/v1/chat/completions",
        headers=headers,
        json=data,
        timeout=600
    )

    response.raise_for_status()  # 触发400错误的异常
    result_data = response.json()
    # print(f"Received API response: {_json.dumps(result_data, indent=2)}")
    if original_chinese_prompt:
        print(f"修改提示词（中文）：{original_chinese_prompt}")
        print(f"修改提示词（英文）：{final_prompt}")
    else:
        print(f"修改提示词：{final_prompt}")

    # 提取生成的图片
    message = result_data.get('choices', [{}])[0].get('message', {})

    if message and 'images' in message and message['images']:
        image_data = message['images'][0]
        base64_string = image_data.get('image_url', {}).get('url', '')

        if base64_string and ',' in base64_string:
            base64_content = base64_string.split(',')[1]
            img_bytes = base64.b64decode(base64_content)
            generated_image = Image.open(_io.BytesIO(img_bytes))
        else:
            raise RuntimeError(f"API返回无效图片格式: {_json.dumps(result_data, indent=2)}")
    else:
        raise RuntimeError(f"API未返回图片: {_json.dumps(result_data, indent=2)}")

    log_detail: str = message.get('content', '无法生成图片')
    
    # 返回更多信息用于保存txt文件
    return generated_image, log_detail, original_chinese_prompt, final_prompt, image_cache_paths


def generate_by_nano_banana_entry(prompt: str, image1, image2, image3, api_key, save_path):
    generated_img = None
    log_detail = "生成图片失败"
    try:
        generated_img, log_detail, original_chinese_prompt, final_prompt, image_cache_paths = generate_by_nano_banana(prompt, image1, image2, image3, api_key)
        
        # 保存生成的图片
        if generated_img and save_path:
            # 确保保存目录存在
            os.makedirs(save_path, exist_ok=True)
            
            # 生成时间戳文件名
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"OR_{timestamp}.png"
            full_path = os.path.join(save_path, filename)
            
            # 保存图片
            generated_img.save(full_path)
            
            # 保存同名txt文件
            txt_filename = f"OR_{timestamp}.txt"
            txt_full_path = os.path.join(save_path, txt_filename)
            
            try:
                with open(txt_full_path, 'w', encoding='utf-8') as f:
                    f.write("=== 图片生成信息 ===\n\n")
                    
                    # 1. 输入图片的本地缓存路径
                    f.write("输入图片本地缓存路径：\n")
                    if image_cache_paths:
                        for i, cache_path in enumerate(image_cache_paths):
                            f.write(f"图片{i+1}: {os.path.abspath(cache_path)}\n")
                    else:
                        f.write("无输入图片\n")
                    
                    f.write("\n")
                    
                    # 2. 中文和英文提示词
                    if original_chinese_prompt:
                        f.write(f"中文提示词：\n{original_chinese_prompt}\n\n")
                        f.write(f"英文提示词：\n{final_prompt}\n")
                    else:
                        f.write(f"提示词：\n{final_prompt}\n")
                        
                print(f"信息文件已保存至: {txt_full_path}")
            except Exception as txt_e:
                print(f"保存txt文件失败: {txt_e}")
            
            # 更新日志显示信息
            if original_chinese_prompt:
                log_detail += f"\n恭喜您，图片已成功生成！\n\n输入的中文提示词：{original_chinese_prompt}"
            else:
                log_detail += f"\n恭喜您，图片已成功生成！"
            
    except Exception as e:
        print(f"OpenRouter错误: {str(e)}")
        log_detail = f"错误: {str(e)}"

    return generated_img, log_detail


with gr.Blocks() as demo:
    # 加载样式数据
    style_names, styles_dict = load_styles()
    
    gr.HTML("""
    <div style='text-align: center; margin-bottom: 30px; padding: 20px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 10px; color: white;'><h1 style='font-size: 2.5rem; margin: 0; text-shadow: 2px 2px 4px rgba(0,0,0,0.3);'>🎨 Nano Banana V2.1.0  </h1><p style='font-size: 1.2rem; margin: 10px 0 0 0; opacity: 0.9;'>用Nano Banana释放你的想象力！Design By：BOZOYAN </p>
    </div>
    """)

    with gr.Row():
        with gr.Column():
            # 添加样式选择功能
            with gr.Row():
                style_dropdown = gr.Dropdown(
                    choices=style_names,
                    label="第一步：选择提示词样式",
                    value=None,
                    interactive=True
                )
                preview_btn = gr.Button("第二步：提示词效果查看", variant="secondary")
            
            prompt = gr.Textbox(label="输入修改提示词：默认不输入是手办玩偶的英文提示词", lines=8,
                value="create a 1/7 scale commercialized figure of thecharacter in the illustration, in a realistic styie and environment.Place the figure on a computer desk, using a circular transparent acrylic base without any text.On the computer screen, display the ZBrush modeling process of the figure.Next to the computer screen, place a BANDAl-style toy packaging box printedwith the original artwork.")
            init_image_1 = gr.Image(label="第三步：输入第一张图", type='pil')
            init_image_2 = gr.Image(label="可选-输入第二张图", type='pil')
            init_image_3 = gr.Image(label="可选-输入第三张图", type='pil')

        with gr.Column():
            # 添加API Key输入框和查询按钮
            with gr.Row():
                openrouter_key_input = gr.Textbox(
                    label="申请地址：https://openrouter.ai/，邮箱免费注册,免费申请一个API Key再输入。", 
                    type="password",
                    placeholder="输入API Key（优先级：最好是自定义输入KEY）",
                    scale=4
                )
                query_key_btn = gr.Button("查询Key使用情况", variant="secondary", scale=1)
            
            # 添加图片保存路径输入框
            default_download_path = get_default_download_path()
            save_path_input = gr.Textbox(
                label="图片保存路径",
                value=default_download_path,
                # visible=False,
                placeholder="输入图片保存的文件夹路径"
            )           
            log_detail_box = gr.Textbox(label="响应日志", interactive=False, lines=4) 
            output_image = gr.Image(label="生成的图片", type='pil', image_mode='RGB', format='png')
            generate_btn = gr.Button("第四步：生成图片", variant="primary")

    generate_btn.click(
        fn=generate_by_nano_banana_entry,
        inputs=[prompt, init_image_1, init_image_2, init_image_3, openrouter_key_input, save_path_input],
        outputs=[output_image, log_detail_box],
    )
    
    # 效果查看按钮事件
    preview_btn.click(
        fn=lambda selected_style: preview_style(selected_style, styles_dict),
        inputs=[style_dropdown],
        outputs=[prompt, log_detail_box, output_image]
    )
    
    # 查询API Key使用情况按钮事件
    query_key_btn.click(
        fn=query_api_key_usage,
        inputs=[openrouter_key_input],
        outputs=[log_detail_box]
    )

# 设置并获取可用端口
port = setup_gradio_url()

demo.queue(default_concurrency_limit=10).launch(
    inbrowser=True,
    server_name="0.0.0.0",  # 允许外部访问
    server_port=port,
    share=False  # 可以设置为True来获取公共URL
)
